package org.mq.mymq.broker.rebalance;

import com.alibaba.fastjson2.JSON;
import org.apache.commons.collections4.CollectionUtils;
import org.mq.mymq.broker.cache.CommonCache;
import org.mq.mymq.broker.model.MyMqTopicModel;
import org.mq.mymq.broker.rebalance.strategy.IReBalanceStrategy;
import org.mq.mymq.broker.rebalance.strategy.ReBalanceInfo;
import org.mq.mymq.broker.rebalance.strategy.impl.RandomReBalanceStrategyImpl;
import org.mq.mymq.broker.rebalance.strategy.impl.RangeReBalanceStrategyImpl;
import org.mq.mymq.common.utils.AssertUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * @ClassName ConsumerInstancePool
 * @Description
 * @Author jiarong_ye
 * @Date 2025/6/28 13:13
 * @Version 1.0
 */
public class ConsumerInstancePool {
    private static final Logger logger = LoggerFactory.getLogger(ConsumerInstancePool.class);
    private Map<String, List<ConsumerInstance>> consumerInstanceMap = new ConcurrentHashMap<>();
    private static Map<String, IReBalanceStrategy> reBalanceStrategyMap = new HashMap<>();
    private ReBalanceInfo reBalanceInfo = new ReBalanceInfo();

    static {
        reBalanceStrategyMap.put("random", new RandomReBalanceStrategyImpl());
        reBalanceStrategyMap.put("range", new RangeReBalanceStrategyImpl());
    }

    /**
     * 加入消费者实例池中
     *
     * @param consumerInstance
     */
    public void addInstancePool(ConsumerInstance consumerInstance) {
        synchronized (this) {
            String topic = consumerInstance.getTopic();
            MyMqTopicModel myMqTopicModel = CommonCache.getMyMqTopicModelsMap().get(topic);
            AssertUtil.isNotNull(myMqTopicModel, "topic非法");
            List<ConsumerInstance> consumerInstances = consumerInstanceMap.getOrDefault(topic, new ArrayList<>());
            for (ConsumerInstance instance : consumerInstances) {
                if (instance.getConsumerRegId().equals(consumerInstance.getConsumerRegId())) {
                    return;
                }
            }
            consumerInstances.add(consumerInstance);
            consumerInstanceMap.put(topic, consumerInstances);
            Set<String> consumerGroupSet = reBalanceInfo.getChangeConsumerGroupMap().get(topic);
            if (CollectionUtils.isEmpty(consumerGroupSet)) {
                consumerGroupSet = new HashSet<>();
            }
            consumerGroupSet.add(consumerInstance.getConsumeGroup());
            reBalanceInfo.getChangeConsumerGroupMap().put(topic, consumerGroupSet);
            logger.info("new instance add in poll:{}", JSON.toJSONString(consumerInstance));
        }
    }

    /**
     * 执行重平衡逻辑
     */
    public void doRebalance() {
        synchronized (this) {
            String reBalanceStrategy = CommonCache.getGlobalProperties().getReBalanceStrategy();
            //触发重平衡行为，根据参数决定重平衡策略的不同
            reBalanceInfo.setConsumeInstanceMap(this.consumerInstanceMap);
            reBalanceStrategyMap.get(reBalanceStrategy).doReBalance(reBalanceInfo);
            reBalanceInfo.getChangeConsumerGroupMap().clear();
        }
    }

    public void startReBalanceJob() {
        Thread reBalanceJob = new Thread(() -> {
           while (true) {
               try {
                   TimeUnit.SECONDS.sleep(10);
                   doRebalance();
               } catch (InterruptedException e) {
                   logger.error("do rebalance error", e);
               }
           }
        });
        reBalanceJob.setName("reBalance-task");
        reBalanceJob.start();
    }

    /**
     * 从池中删除消费者实例
     *
     * @param reqId
     */
    public void removeInstancePool(String reqId) {
        synchronized (this) {
            for (String topic : consumerInstanceMap.keySet()) {
                List<ConsumerInstance> consumerInstances = consumerInstanceMap.get(topic);
                List<ConsumerInstance> filterInstances = consumerInstances.stream().filter(item -> !item.getConsumerRegId().equals(reqId)).collect(Collectors.toList());
                //剔除断开的消费者节点
                consumerInstanceMap.put(topic, filterInstances);
            }
            Map<String, Map<String, List<ConsumerInstance>>> consumeHoldMap = CommonCache.getConsumerHoldMap();
            for (String topic : consumeHoldMap.keySet()) {
                Map<String,List<ConsumerInstance>> consumeGroupInstanceMap = consumeHoldMap.get(topic);
                for (String consumeGroup : consumeGroupInstanceMap.keySet()) {
                    List<ConsumerInstance> consumerInstances = consumeGroupInstanceMap.get(consumeGroup);
                    List<ConsumerInstance> filterInstances = consumerInstances.stream().filter(item -> !item.getConsumerRegId().equals(reqId)).collect(Collectors.toList());
                    consumeGroupInstanceMap.put(consumeGroup,filterInstances);
                }
            }
        }
    }
}
